from PyQt5.QtWidgets import (QMainWindow, QApplication, QPushButton, QGridLayout, QWidget, QLabel, QTextEdit, QHBoxLayout, QVBoxLayout,
                             QDesktopWidget, QListWidget, QLineEdit, QTextBrowser, QComboBox, QMessageBox)
from PyQt5.QtCore import Qt
import sys, os

sys.path.append(os.path.abspath('.'))
sys.path.append(os.path.join(os.path.abspath('.'), 'embedding_model'))
sys.path.append(os.path.join(os.path.abspath('.'), 'amie_plus_model'))

from excavatePossibleTriples import excavatePossibleTriple
# from updateKnowledgeGraph import updateKnowledgeGraph, updateKnowledgeGraphByRules
from updateKnowledgeGraph import updateModel
from config import SystemConfig, AMIEConfig


class System(QMainWindow):

    def __init__(self, config, amie_config):
        super(System, self).__init__()
        self.topk_per_epoch = config.topk_per_epoch
        self.pool_topk = config.pool_topk
        self.k_kdtree = config.k_kdtree
        self.embedding = True
        self.embed_e_triples, self.embed_e_similarities = [], []
        self.rule_e_triples, self.rule_e_similarities = [], []
        self.excavateModel = excavatePossibleTriple(config.data_path, config.dataset, config.model_path, config.gamma,
                                                    config.embedding_model, amie_config)
        # self.embedUpdateModel = updateKnowledgeGraph(config.dataset, config.data_path, config.log_path, config.model_path,
        #                                        config.local_update_threshold, config.global_update_times, config.embedding_model)
        # self.ruleUpdateModel = updateKnowledgeGraphByRules(config.dataset, config.data_path)
        self.updateModel = updateModel(config.dataset, config.data_path, config.log_path, config.model_path,
                                               config.local_update_threshold, config.global_update_times, config.embedding_model)
        print("Start reasoning...")
        self.embed_triples_generator = self.excavateModel.reasoning(config.topk_per_epoch, config.pool_topk, config.k_kdtree,
                                                              True)
        self.rule_triples_generator = self.excavateModel.reasoning(config.topk_per_epoch, config.pool_topk, config.k_kdtree,
                                                              False)

        self.exist_pos = 0
        self.update = False
        self.initUI()

    def initUI(self):
        self.setFixedSize(1000, 540)
        self.center()
        self.setWindowTitle("Anti-QA")

        self.widget = QWidget(self)
        self.setCentralWidget(self.widget)

        self.editWidget()
        self.state = None
        # self.rule_e_triples, self.rule_e_similarities, _ = next(self.triples_generator)
        # self.embed_e_triples = ["起始状态\t对电网进行简化等值\t电网无死岛",
        #                   "计算不收敛\t操作\t分区之间有联络线",
        #                   "计算不收敛\t操作\t更改计算方法",
        #                   "该分区各电压等级下无功发电大于无功负荷\t操作\t投入并联电抗器",
        #                   "母线\t有属性\t基准电压",
        #                   "有变压器与最大误差母线相连\t操作\t更改计算精度",
        #                   "减少发电机出力\t检查发电机有功功率是否达到下限\t预设平衡点",
        #                   "分区有联络线\t操作\t检查有功功率是否过大",
        #                   "电网规模较大\t操作\t更改初值",
        #                   "有变压器与最大误差母线相连\t操作\t检查该变压器阻抗和变比"
        #                   ]
        # self.embed_e_similarities = [0.2385722, 0.2385722, 0.2385722, 0.2284294, 0.2284294,
        #                        0.2282942, 0.2198329, 0.2198329, 0.2198329, 0.2030239]

        print("Finish reasoning...")
        # self.addListItems()

        self.show()

    def editWidget(self):
        self.modeLabel = QLabel("模式: ", self)
        self.mode_choice = QComboBox(self)
        # self.mode_choice.addItems(['Embedding 推理'])
        self.mode_choice.addItems(['Embed 推理', '规则推理'])

        self.embedList = QListWidget(self)
        self.embedListLabel = QLabel("可能存在的知识", self)
        self.generateBtn = QPushButton("生成", self)

        self.existsList = QListWidget(self)
        self.existsListLabel = QLabel("存在的知识", self)
        self.inexistentList = QListWidget(self)
        self.inexistentListLabel = QLabel("不存在的知识", self)
        self.exist_clean_btn = QPushButton("清除")
        self.inexist_clean_btn = QPushButton("清除")

        self.triple_label = QLabel("三元组: ")
        self.possibility = QLabel("可能性: ")
        self.space_label = QLabel("")
        self.lineEdit = QTextBrowser(self)
        self.lineEdit.setDisabled(True)
        self.okbtn = QPushButton("yes", self)
        self.cancelbtn = QPushButton("no", self)

        self.editAction()
        self.editLayout()

    def editLayout(self):
        hbox = QHBoxLayout()

        right_vbox = QVBoxLayout()
        middle_vbox = QVBoxLayout()
        left_vbox = QVBoxLayout()
        left_vbox_up = QHBoxLayout()
        left_vbox_down = QVBoxLayout()

        left_vbox_up.addWidget(self.modeLabel)
        left_vbox_up.addWidget(self.mode_choice)

        grid = QGridLayout()
        grid.setSpacing(20)
        grid.addWidget(self.embedListLabel, 0, 0)
        grid.addWidget(self.generateBtn, 0, 2)
        grid.addWidget(self.embedList, 1, 0, 20, 3)
        left_vbox_down.addStretch(1)
        left_vbox_down.addLayout(grid)
        left_vbox_down.addStretch(3)

        grid = QGridLayout()
        grid.setSpacing(10)
        grid.addWidget(self.triple_label, 0, 0)
        grid.addWidget(self.lineEdit, 1, 0, 1, 5)
        grid.addWidget(self.possibility, 2, 0)
        grid.addWidget(self.space_label, 2, 1)
        grid.addWidget(self.okbtn, 4, 0)
        grid.addWidget(self.cancelbtn, 4, 4)
        middle_vbox.addStretch(1)
        middle_vbox.addLayout(grid)
        middle_vbox.addStretch(5)

        grid = QGridLayout()
        grid.setSpacing(10)
        grid.addWidget(self.existsListLabel, 0, 0)
        grid.addWidget(self.exist_clean_btn, 0, 1)
        grid.addWidget(self.existsList, 1, 0, 3, 2)
        grid.addWidget(self.inexistentListLabel, 4, 0)
        grid.addWidget(self.inexist_clean_btn, 4, 1)
        grid.addWidget(self.inexistentList, 5, 0, 3, 2)
        right_vbox.addStretch(1)
        right_vbox.addLayout(grid)
        right_vbox.addStretch(5)

        left_vbox.addLayout(left_vbox_up)
        left_vbox.addLayout(left_vbox_down)
        hbox.addLayout(left_vbox)
        hbox.addLayout(middle_vbox)
        hbox.addLayout(right_vbox)

        self.widget.setLayout(hbox)

    def editAction(self):
        self.mode_choice.currentTextChanged.connect(lambda: self.combox_index_change(self.mode_choice.currentText()))

        self.embedList.itemDoubleClicked.connect(self.doubleClickItem)
        self.existsList.itemDoubleClicked.connect(self.doubleClickRightItem)
        self.inexistentList.itemDoubleClicked.connect(self.doubleClickRightItem)

        self.okbtn.clicked.connect(self.clickButton)
        self.cancelbtn.clicked.connect(self.clickButton)
        self.generateBtn.clicked.connect(self.clickGenerateButton)
        self.exist_clean_btn.clicked.connect(self.cleanExistList)
        self.inexist_clean_btn.clicked.connect(self.cleanInExistList)

    def addListItems(self):
        # contents = []
        # count = self.embedList.count()
        # # 遍历listwidget中的内容
        # for i in range(count):
        #     contents.append(self.embedList.item(i).text())

        if self.embedding:
            triples, similarities = self.embed_e_triples, self.embed_e_similarities
        else:
            triples, similarities = self.rule_e_triples, self.rule_e_similarities

        # self.embedList.clear()
        for e_triple, e_similarity in zip(triples, similarities):
            e_piece = str(e_triple) + "\t" + str(e_similarity)
            self.embedList.addItem(e_piece)
        if triples:
            text = "\t".join([triples[0], str(similarities[0])])
            self.showTripleAndPossibility(text)

    def combox_index_change(self, choice):
        if choice == 'Embed 推理' and not self.embedding:
            self.embedding = True
            self.embedList.clear()
            self.addListItems()
            print("embedding")
        elif choice == '规则推理' and self.embedding:
            self.embedding = False
            self.embedList.clear()
            self.addListItems()
            print("rule")

    def doubleClickItem(self, item):
        line = item.text()
        self.showTripleAndPossibility(line)

    def doubleClickRightItem(self, item):
        self.doubleClickItem(item)
        text = item.text()
        find_item = self.existsList.findItems(text, Qt.MatchContains)
        if find_item:
            row = self.existsList.row(find_item[0])
            self.existsList.takeItem(row)
        else:
            find_item = self.inexistentList.findItems(text, Qt.MatchContains)
            if find_item:
                row = self.inexistentList.row(find_item[0])
                self.inexistentList.takeItem(row)
        self.embedList.insertItem(0, item)

    def clickButton(self):
        sender = self.sender()
        text = self.lineEdit.toPlainText()
        possibility = self.space_label.text()
        if text:
            text = "\t".join(text.split("\n")) + "\t" + possibility
            if sender.text() == "yes":
                if text != "":
                    self.existsList.addItem(text)
            elif sender.text() == "no":
                if text != "":
                    self.inexistentList.addItem(text)
            # flag = True
            item = self.embedList.findItems(text, Qt.MatchContains)
            if item:
                row = self.embedList.row(item[0])
                self.embedList.takeItem(row)

            if self.embedList.count() > 0:
                text = self.embedList.item(0).text()
                self.showTripleAndPossibility(text)
            else:
                self.lineEdit.setText("")
                self.space_label.setText("")
                try:
                    self.e_triples, self.e_similarities, remain_triples = None, None, 0
                    print("remain: ", remain_triples)
                    # self.addListItems()
                except:
                    pass

    def cleanExistList(self):
        self.existsList.clear()
        self.exist_pos = 0

    def cleanInExistList(self):
        self.inexistentList.clear()

    def showTripleAndPossibility(self, text):
        elems = text.split("\t")
        triple, possibility = elems[:3], elems[3]
        triple = "\n".join(triple)
        self.lineEdit.setText(triple)
        self.space_label.setText(possibility)

    def clickGenerateButton(self):
        if not self.update:
            self.update = True
            triples = []
            for pos in range(self.exist_pos, self.existsList.count()):
                item = self.existsList.item(pos).text()
                item = item.split("\t")[:3]
                triples.append(item)
            self.exist_pos = self.existsList.count()
            print("开始更新...")
            if len(triples) > 0:
                if self.embedding:
                    self.updateModel.updateGraphEmbedding(triples)
                else:
                    self.updateModel.updateKnowledgeGraphByRules(triples)
                    # pass
            print("开始挖掘")
            self.excavateTriples()
            self.update = False
            print("更新结束")

    def excavateTriples(self):
        try:
            if self.embedding:
                self.embed_e_triples, self.embed_e_similarities, remain_triples = next(self.embed_triples_generator)
                # self.embed_e_triples, self.embed_e_similarities, remain_triples = None, None, 0
            else:
                self.rule_e_triples, self.rule_e_similarities, remain_triples = next(self.rule_triples_generator)
                if not self.rule_e_triples:
                    QMessageBox.warning(self, 'Warning', '未挖掘出规则', QMessageBox.Yes)
                    return
                # self.rule_e_triples, self.rule_e_similarities, remain_triples = None, None, 0
        except:
            if self.embedding:
                self.embed_triples_generator = self.excavateModel.reasoning(config.topk_per_epoch, config.pool_topk,
                                                                              config.k_kdtree,
                                                                              True)
                self.embed_e_triples, self.embed_e_similarities, remain_triples = next(self.embed_triples_generator)
                # self.embed_e_triples, self.embed_e_similarities, remain_triples = None, None, 0
            else:
                self.rule_triples_generator = self.excavateModel.reasoning(config.topk_per_epoch, config.pool_topk,
                                                                           config.k_kdtree,
                                                                           False)
                self.rule_e_triples, self.rule_e_similarities, remain_triples = next(self.rule_triples_generator)
                if not self.rule_e_triples:
                    QMessageBox.warning(self, 'Warning', '未挖掘出规则', QMessageBox.Yes)
                    return
        print("remain: ", remain_triples)
        self.addListItems()

    def center(self):
        qr = self.frameGeometry()

        cp = QDesktopWidget().availableGeometry().center()
        qr.moveCenter(cp)
        self.move(qr.topLeft())



if __name__ == "__main__":
    dataset = "PowerGrid"
    data_path = os.path.join("..", "data", dataset)
    use_delta = False
    gamma, model_name, topK, pool_topk, k_kdtree = 24.0, "TransE", 10, 30, 5
    local_update_threshold = 2              # 累计新增知识超过该值将进行局部更新，更新后累计新增知识数清0
    global_update_times = 2                 # 累计局部更新次数超过该值将全局更新，更新后累计局部更新次数清0
    log_path = os.path.join("..", "record", dataset)
    model_path = os.path.join("..", "model", dataset)
    config = SystemConfig(
        dataset=dataset,
        data_path=data_path,
        model_path=model_path,
        log_path=log_path,
        gamma=gamma,
        embedding_model=model_name,
        topk_per_epoch=topK,
        pool_topk=pool_topk,
        k_kdtree=k_kdtree,
        global_update_times=global_update_times,
        local_update_threshold=local_update_threshold
    )

    max_rule_len, min_pcaconf, num_process, result_dir = 3, 0.8, 4, os.path.join("..", "result", dataset, "amie")
    amie_config = AMIEConfig(dataset=dataset, data_dir=data_path, maxad=max_rule_len, minpca=min_pcaconf,
                             num_process=num_process, result_dir=result_dir)

    app = QApplication(sys.argv)
    system = System(config, amie_config)
    # ex = Example()
    sys.exit(app.exec_())